from collections.abc import AsyncGenerator
from typing import TYPE_CHECKING, Any, Optional
from uuid import UUID

from posthog.schema import AgentMode, AssistantMessage, HumanMessage, MaxBillingContext

from posthog.models import Team, User

from ee.hogai.chat_agent import AssistantGraph
from ee.hogai.chat_agent.stream_processor import ChatAgentStreamProcessor
from ee.hogai.chat_agent.taxonomy.types import TaxonomyNodeName
from ee.hogai.core.runner import BaseAgentRunner
from ee.hogai.utils.types import (
    AssistantMode,
    AssistantNodeName,
    AssistantOutput,
    AssistantState,
    PartialAssistantState,
)
from ee.models import Conversation

if TYPE_CHECKING:
    from ee.hogai.utils.types.composed import MaxNodeName


STREAMING_NODES: set["MaxNodeName"] = {
    AssistantNodeName.ROOT,
    AssistantNodeName.INKEEP_DOCS,
    AssistantNodeName.MEMORY_ONBOARDING,
    AssistantNodeName.MEMORY_INITIALIZER,
    AssistantNodeName.MEMORY_ONBOARDING_ENQUIRY,
    AssistantNodeName.MEMORY_ONBOARDING_FINALIZE,
    AssistantNodeName.DASHBOARD_CREATION,
}


VERBOSE_NODES: set["MaxNodeName"] = {
    AssistantNodeName.TRENDS_GENERATOR,
    AssistantNodeName.FUNNEL_GENERATOR,
    AssistantNodeName.RETENTION_GENERATOR,
    AssistantNodeName.SQL_GENERATOR,
    AssistantNodeName.INSIGHTS_SEARCH,
    AssistantNodeName.QUERY_EXECUTOR,
    AssistantNodeName.MEMORY_INITIALIZER_INTERRUPT,
    AssistantNodeName.ROOT,
    AssistantNodeName.ROOT_TOOLS,
    AssistantNodeName.SLASH_COMMAND_HANDLER,
    TaxonomyNodeName.TOOLS_NODE,
    TaxonomyNodeName.TASK_EXECUTOR,
}


class ChatAgentRunner(BaseAgentRunner):
    _state: Optional[AssistantState]
    _initial_state: Optional[AssistantState | PartialAssistantState]
    _selected_agent_mode: AgentMode | None

    def __init__(
        self,
        team: Team,
        conversation: Conversation,
        *,
        new_message: Optional[HumanMessage] = None,
        user: User,
        session_id: Optional[str] = None,
        contextual_tools: Optional[dict[str, Any]] = None,
        is_new_conversation: bool = False,
        trace_id: Optional[str | UUID] = None,
        billing_context: Optional[MaxBillingContext] = None,
        initial_state: Optional[AssistantState | PartialAssistantState] = None,
        agent_mode: AgentMode | None = None,
    ):
        super().__init__(
            team,
            conversation,
            new_message=new_message,
            user=user,
            graph=AssistantGraph(team, user).compile_full_graph(),
            state_type=AssistantState,
            partial_state_type=PartialAssistantState,
            mode=AssistantMode.ASSISTANT,
            session_id=session_id,
            contextual_tools=contextual_tools,
            is_new_conversation=is_new_conversation,
            trace_id=trace_id,
            billing_context=billing_context,
            stream_processor=ChatAgentStreamProcessor(
                verbose_nodes=VERBOSE_NODES,
                streaming_nodes=STREAMING_NODES,
                state_type=AssistantState,
                team=team,
                user=user,
            ),
        )
        self._selected_agent_mode = agent_mode

    def get_initial_state(self) -> AssistantState:
        if self._latest_message:
            new_state = AssistantState(
                messages=[self._latest_message],
                start_id=self._latest_message.id,
                query_generation_retry_count=0,
                graph_status=None,
                rag_context=None,
            )
            # Only set the agent mode if it was explicitly set.
            if self._selected_agent_mode:
                new_state.agent_mode = self._selected_agent_mode
            return new_state

        # When resuming, do not set the mode. It should start from the same mode as the previous generation.
        return AssistantState(messages=[])

    def get_resumed_state(self) -> PartialAssistantState:
        if not self._latest_message:
            return PartialAssistantState(messages=[])
        new_state = PartialAssistantState(
            messages=[self._latest_message], graph_status="resumed", query_generation_retry_count=0
        )
        # Only set the agent mode if it was explicitly set.
        if self._selected_agent_mode:
            new_state.agent_mode = self._selected_agent_mode
        return new_state

    async def astream(
        self,
        stream_message_chunks: bool = True,
        stream_subgraphs: bool = True,
        stream_first_message: bool = True,
        stream_only_assistant_messages: bool = False,
    ) -> AsyncGenerator[AssistantOutput, None]:
        last_ai_message: AssistantMessage | None = None
        async for stream_event in super().astream(
            stream_message_chunks, stream_subgraphs, stream_first_message, stream_only_assistant_messages
        ):
            _, message = stream_event
            if isinstance(message, AssistantMessage):
                last_ai_message = message
            yield stream_event

        output = last_ai_message.content if isinstance(last_ai_message, AssistantMessage) else None
        await self._report_conversation_state(
            "chat with ai",
            {
                "prompt": self._latest_message.content if self._latest_message else None,
                "output": output,
                "is_new_conversation": self._is_new_conversation,
                "$session_id": self._session_id,
            },
        )
